from docx import Document
from docx.document import Document as DocumentType
from docx.table import Table
from docx.text.paragraph import Paragraph
from docx.shared import Inches
from typing import Dict, Any, List, Optional
from io import BytesIO
import re

class DOCXProcessor:
    """DOCX文档处理器"""
    
    def __init__(self):
        self.supported_formats = ['docx']
    
    def extract_content(self, file_content: bytes) -> Dict[str, Any]:
        """
        从DOCX文件中提取内容和结构信息
        
        Args:
            file_content: DOCX文件的字节内容
            
        Returns:
            Dict[str, Any]: 包含文本内容和结构信息的字典
        """
        
        try:
            # 从字节流创建Document对象
            doc = Document(BytesIO(file_content))
            
            # 提取基本信息
            metadata = self._extract_metadata(doc)
            
            # 提取文本内容
            text_content = self._extract_text(doc)
            
            # 提取结构信息
            structure = self._extract_structure(doc)
            
            # 提取表格信息
            tables = self._extract_tables(doc)
            
            # 提取图片信息
            images = self._extract_images(doc)
            
            # 提取样式信息
            styles = self._extract_styles(doc)
            
            result = {
                'text': text_content,
                'metadata': metadata,
                'structure': structure,
                'tables': tables,
                'images': images,
                'styles': styles,
                'format': 'docx'
            }
            
            return result
            
        except Exception as e:
            raise Exception(f"DOCX处理失败: {str(e)}")
    
    def _extract_metadata(self, doc: DocumentType) -> Dict[str, Any]:
        """提取DOCX元数据"""
        
        core_props = doc.core_properties
        
        metadata = {
            'title': core_props.title or '',
            'author': core_props.author or '',
            'subject': core_props.subject or '',
            'keywords': core_props.keywords or '',
            'comments': core_props.comments or '',
            'category': core_props.category or '',
            'created': str(core_props.created) if core_props.created else '',
            'modified': str(core_props.modified) if core_props.modified else '',
            'last_modified_by': core_props.last_modified_by or '',
            'revision': core_props.revision or 0,
            'version': core_props.version or '',
            'paragraph_count': len(doc.paragraphs),
            'table_count': len(doc.tables)
        }
        
        return metadata
    
    def _extract_text(self, doc: DocumentType) -> str:
        """提取DOCX文本内容"""
        
        text_parts = []
        
        # 遍历所有段落和表格，按文档顺序
        for element in doc.element.body:
            if element.tag.endswith('p'):  # 段落
                # 找到对应的段落对象
                for para in doc.paragraphs:
                    if para._element == element:
                        para_text = self._extract_paragraph_text(para)
                        if para_text.strip():
                            text_parts.append(para_text)
                        break
            
            elif element.tag.endswith('tbl'):  # 表格
                # 找到对应的表格对象
                for table in doc.tables:
                    if table._element == element:
                        table_text = self._extract_table_text(table)
                        if table_text.strip():
                            text_parts.append(f"\n[表格开始]\n{table_text}\n[表格结束]\n")
                        break
        
        return '\n'.join(text_parts)
    
    def _extract_paragraph_text(self, paragraph: Paragraph) -> str:
        """提取段落文本"""
        
        try:
            # 获取段落的完整文本
            text = paragraph.text
            
            # 检查段落样式
            style_name = paragraph.style.name if paragraph.style else ''
            
            # 为标题添加标记
            if 'Heading' in style_name or 'Title' in style_name:
                level = self._extract_heading_level(style_name)
                text = f"{'#' * level} {text}"
            
            return text
            
        except Exception:
            return paragraph.text if paragraph.text else ''
    
    def _extract_heading_level(self, style_name: str) -> int:
        """从样式名称提取标题级别"""
        
        # 匹配 Heading 1, Heading 2 等
        match = re.search(r'Heading\s*(\d+)', style_name)
        if match:
            return int(match.group(1))
        
        # 其他标题样式
        if 'Title' in style_name:
            return 1
        elif 'Subtitle' in style_name:
            return 2
        
        return 1
    
    def _extract_table_text(self, table: Table) -> str:
        """提取表格文本"""
        
        try:
            table_data = []
            
            for row in table.rows:
                row_data = []
                for cell in row.cells:
                    cell_text = cell.text.strip().replace('\n', ' ')
                    row_data.append(cell_text)
                
                if any(cell.strip() for cell in row_data):  # 跳过空行
                    table_data.append(' | '.join(row_data))
            
            return '\n'.join(table_data)
            
        except Exception:
            return '[表格解析失败]'
    
    def _extract_structure(self, doc: DocumentType) -> Dict[str, Any]:
        """提取DOCX结构信息"""
        
        structure = {
            'headings': [],
            'sections': [],
            'paragraphs': [],
            'lists': []
        }
        
        # 提取标题结构
        for i, para in enumerate(doc.paragraphs):
            style_name = para.style.name if para.style else ''
            text = para.text.strip()
            
            if text:
                para_info = {
                    'index': i,
                    'text': text,
                    'style': style_name,
                    'is_heading': 'Heading' in style_name or 'Title' in style_name
                }
                
                structure['paragraphs'].append(para_info)
                
                # 如果是标题，添加到标题列表
                if para_info['is_heading']:
                    level = self._extract_heading_level(style_name)
                    structure['headings'].append({
                        'text': text,
                        'level': level,
                        'style': style_name,
                        'paragraph_index': i
                    })
        
        # 提取列表信息
        structure['lists'] = self._extract_lists(doc)
        
        # 根据标题生成章节结构
        structure['sections'] = self._build_sections(structure['headings'])
        
        return structure
    
    def _extract_lists(self, doc: DocumentType) -> List[Dict[str, Any]]:
        """提取列表信息"""
        
        lists = []
        current_list = None
        
        for i, para in enumerate(doc.paragraphs):
            style_name = para.style.name if para.style else ''
            text = para.text.strip()
            
            # 检查是否是列表项
            if self._is_list_paragraph(para, style_name):
                if current_list is None:
                    current_list = {
                        'start_index': i,
                        'items': [],
                        'type': self._get_list_type(style_name)
                    }
                
                current_list['items'].append({
                    'text': text,
                    'paragraph_index': i,
                    'style': style_name
                })
            
            else:
                # 如果当前段落不是列表项，结束当前列表
                if current_list is not None:
                    current_list['end_index'] = i - 1
                    lists.append(current_list)
                    current_list = None
        
        # 处理文档末尾的列表
        if current_list is not None:
            current_list['end_index'] = len(doc.paragraphs) - 1
            lists.append(current_list)
        
        return lists
    
    def _is_list_paragraph(self, para: Paragraph, style_name: str) -> bool:
        """判断段落是否是列表项"""
        
        # 检查样式名称
        list_styles = ['List', 'Bullet', 'Number']
        if any(style in style_name for style in list_styles):
            return True
        
        # 检查文本开头
        text = para.text.strip()
        if not text:
            return False
        
        # 检查是否以列表标记开头
        list_patterns = [
            r'^[•·▪▫◦‣⁃]\s+',  # 项目符号
            r'^\d+[.).]\s+',    # 数字列表
            r'^[a-zA-Z][.).]\s+',  # 字母列表
            r'^[ivxlcdm]+[.).]\s+',  # 罗马数字
            r'^[-*+]\s+'        # 破折号列表
        ]
        
        for pattern in list_patterns:
            if re.match(pattern, text, re.IGNORECASE):
                return True
        
        return False
    
    def _get_list_type(self, style_name: str) -> str:
        """获取列表类型"""
        
        if 'Bullet' in style_name:
            return 'bullet'
        elif 'Number' in style_name:
            return 'numbered'
        else:
            return 'unknown'
    
    def _build_sections(self, headings: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """根据标题构建章节结构"""
        
        sections = []
        
        for i, heading in enumerate(headings):
            section = {
                'title': heading['text'],
                'level': heading['level'],
                'start_paragraph': heading['paragraph_index'],
                'subsections': []
            }
            
            # 找到下一个同级或更高级标题的位置
            end_paragraph = None
            for j in range(i + 1, len(headings)):
                next_heading = headings[j]
                if next_heading['level'] <= heading['level']:
                    end_paragraph = next_heading['paragraph_index'] - 1
                    break
            
            section['end_paragraph'] = end_paragraph
            sections.append(section)
        
        return sections
    
    def _extract_tables(self, doc: DocumentType) -> List[Dict[str, Any]]:
        """提取表格信息"""
        
        tables = []
        
        for i, table in enumerate(doc.tables):
            try:
                table_info = {
                    'index': i,
                    'rows': len(table.rows),
                    'columns': len(table.columns) if table.rows else 0,
                    'data': [],
                    'headers': []
                }
                
                # 提取表格数据
                for row_idx, row in enumerate(table.rows):
                    row_data = []
                    for cell in row.cells:
                        cell_text = cell.text.strip()
                        row_data.append(cell_text)
                    
                    table_info['data'].append(row_data)
                    
                    # 假设第一行是表头
                    if row_idx == 0:
                        table_info['headers'] = row_data
                
                # 计算表格统计信息
                table_info['total_cells'] = sum(len(row) for row in table_info['data'])
                table_info['non_empty_cells'] = sum(
                    1 for row in table_info['data'] 
                    for cell in row if cell.strip()
                )
                
                tables.append(table_info)
                
            except Exception as e:
                tables.append({
                    'index': i,
                    'error': f'表格解析失败: {str(e)}'
                })
        
        return tables
    
    def _extract_images(self, doc: DocumentType) -> List[Dict[str, Any]]:
        """提取图片信息"""
        
        images = []
        
        try:
            # 遍历文档中的所有关系
            for rel in doc.part.rels.values():
                if "image" in rel.target_ref:
                    try:
                        image_info = {
                            'relationship_id': rel.rId,
                            'target': rel.target_ref,
                            'type': rel.reltype,
                            'size': len(rel.target_part.blob) if hasattr(rel.target_part, 'blob') else 0
                        }
                        
                        # 尝试获取图片的更多信息
                        if hasattr(rel.target_part, 'blob'):
                            blob = rel.target_part.blob
                            image_info['format'] = self._detect_image_format(blob)
                        
                        images.append(image_info)
                        
                    except Exception as e:
                        images.append({
                            'relationship_id': rel.rId,
                            'error': f'图片信息提取失败: {str(e)}'
                        })
        
        except Exception:
            pass
        
        return images
    
    def _detect_image_format(self, blob: bytes) -> str:
        """检测图片格式"""
        
        if blob.startswith(b'\xff\xd8\xff'):
            return 'jpeg'
        elif blob.startswith(b'\x89PNG\r\n\x1a\n'):
            return 'png'
        elif blob.startswith(b'GIF87a') or blob.startswith(b'GIF89a'):
            return 'gif'
        elif blob.startswith(b'BM'):
            return 'bmp'
        else:
            return 'unknown'
    
    def _extract_styles(self, doc: DocumentType) -> Dict[str, Any]:
        """提取样式信息"""
        
        styles = {
            'available_styles': [],
            'used_styles': {},
            'style_hierarchy': []
        }
        
        try:
            # 获取所有可用样式
            for style in doc.styles:
                style_info = {
                    'name': style.name,
                    'type': str(style.type),
                    'builtin': style.builtin
                }
                styles['available_styles'].append(style_info)
            
            # 统计使用的样式
            for para in doc.paragraphs:
                style_name = para.style.name if para.style else 'Normal'
                if style_name in styles['used_styles']:
                    styles['used_styles'][style_name] += 1
                else:
                    styles['used_styles'][style_name] = 1
        
        except Exception:
            pass
        
        return styles
    
    def extract_text_with_formatting(self, file_content: bytes) -> List[Dict[str, Any]]:
        """提取带格式信息的文本"""
        
        try:
            doc = Document(BytesIO(file_content))
            
            formatted_text = []
            
            for para_idx, para in enumerate(doc.paragraphs):
                if not para.text.strip():
                    continue
                
                para_info = {
                    'paragraph_index': para_idx,
                    'text': para.text,
                    'style': para.style.name if para.style else 'Normal',
                    'runs': []
                }
                
                # 提取运行（run）级别的格式信息
                for run in para.runs:
                    if run.text:
                        run_info = {
                            'text': run.text,
                            'bold': run.bold,
                            'italic': run.italic,
                            'underline': run.underline,
                            'font_name': run.font.name if run.font.name else '',
                            'font_size': str(run.font.size) if run.font.size else ''
                        }
                        para_info['runs'].append(run_info)
                
                formatted_text.append(para_info)
            
            return formatted_text
            
        except Exception as e:
            raise Exception(f"格式化文本提取失败: {str(e)}")
    
    def extract_comments_and_revisions(self, file_content: bytes) -> Dict[str, Any]:
        """提取批注和修订信息"""
        
        try:
            doc = Document(BytesIO(file_content))
            
            result = {
                'comments': [],
                'revisions': [],
                'has_tracked_changes': False
            }
            
            # 注意：python-docx对批注和修订的支持有限
            # 这里提供基础框架，实际实现可能需要更底层的XML解析
            
            # 检查是否有修订跟踪
            try:
                # 这是一个简化的检查方法
                xml_content = doc._element.xml
                if 'w:ins' in xml_content or 'w:del' in xml_content:
                    result['has_tracked_changes'] = True
            except Exception:
                pass
            
            return result
            
        except Exception as e:
            raise Exception(f"批注和修订提取失败: {str(e)}")